/*! ******************************************************************************
 *
 * Created on 2022年3月28日
 *
 * Copyright(c) YuZhou Big Data Laboratory CO.,Ltd, 2022.  All rights reserved.
 *
 *******************************************************************************
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 ******************************************************************************/

package org.yzbdl.lanius.core.plugins;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.pentaho.di.core.exception.KettlePluginException;
import org.pentaho.di.core.plugins.KettleSelectiveParentFirstClassLoader;
import org.pentaho.di.core.plugins.KettleURLClassLoader;
import org.pentaho.di.core.plugins.Plugin;
import org.pentaho.di.core.plugins.PluginInterface;
import org.pentaho.di.i18n.BaseMessages;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.core.util.URLUtil;

/**
 * @author hs@yzbdl.ac.cn
 *
 */
public class ClassLoaderManager {
    private static final Class<?> PKG = ClassLoaderManager.class;

    private static final ClassLoaderManager clm = new ClassLoaderManager();

    enum CacheType {
        UNIQUE_CODE, USER_DEFINED, FOLDER_BASED
    }

    /**
     * 两级Map, 缓存类加载器. <br/>
     * <br/>
     * 第一级Map: <br/>
     * 记录缓存类型与缓存集合的关联关系, <br/>
     * key:为缓存类型:目前支持uniqueCode、userDefined、folderBased三种 <br/>
     * value:为缓存Map <br/>
     * =================================================<br/>
     * 第二级Map:<br/>
     * 记录缓存key与类加载器的映射.
     */
    private final Map<String, Map<String, URLClassLoader>> classLoaderMap =
        new ConcurrentHashMap<String, Map<String, URLClassLoader>>() {
            private static final long serialVersionUID = 1L;
            {
                put(CacheType.UNIQUE_CODE.name(), Maps.newConcurrentMap());
                put(CacheType.USER_DEFINED.name(), Maps.newConcurrentMap());
                put(CacheType.FOLDER_BASED.name(), Maps.newConcurrentMap());
            }
        };

    /**
     * <pre>
     * 持有独立类加载器的引用, 用于对类加载器进行关闭.
     * 一个插件会关联一个或多个类加载器, 当进行插件卸载时, 需要对插件关联的所有类加载器进行关闭.
     * </pre>
     */
    private final Map<String, List<URLClassLoader>> separateClassLoaderHandler = Maps.newHashMap();

    /**
     * 插件实例与插件中的类，记录一个插件中哪些类通过父类加载器进行加载。
     */
    private final Map<PluginInterface, String[]> parentClassLoaderPatternMap = new HashMap<PluginInterface, String[]>();

    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    private ClassLoaderManager() {}

    public static ClassLoaderManager getInstance() {
        return clm;
    }

    public void addParentClassLoaderPatterns(PluginInterface plugin, String[] patterns) {
        lock.writeLock().lock();
        try {
            parentClassLoaderPatternMap.put(plugin, patterns);
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 获取类加载器，如果没有则创建。
     * @param plugin
     * @return
     * @throws KettlePluginException
     */
    public ClassLoader getClassLoader(PluginInterface plugin) throws KettlePluginException {

        if (plugin == null) {
            throw new KettlePluginException(
                BaseMessages.getString(PKG, "PluginRegistry.RuntimeError.NoValidStepOrPlugin.PLUGINREGISTRY001"));
        }

        // 本地插件, 使用主程序的类加载器?
        if (plugin.isNativePlugin()) {
            return this.getClass().getClassLoader();
        } else {
            URLClassLoader ucl = null;
            lock.writeLock().lock();
            try {
                if (plugin.isSeparateClassLoaderNeeded()) {
                    ucl = createClassLoader(plugin);

                    // 记录到separateClassLoaderHandler
                    String cacheKey = ((Plugin)plugin).getPluginUniqueCode();
                    List<URLClassLoader> classLoaders =
                        separateClassLoaderHandler.getOrDefault(cacheKey, Lists.newArrayList());
                    classLoaders.add(ucl);

                } else {
                    ucl = getCachedClassLoader(plugin);
                }
            } finally {
                lock.writeLock().unlock();
            }
            return ucl;
        }
    }

    public void removeClassLoader(PluginInterface plugin) {
        // if (plugin == null) {
        // throw new KettlePluginException(
        // BaseMessages.getString(PKG, "PluginRegistry.RuntimeError.NoValidStepOrPlugin.PLUGINREGISTRY001"));
        // }
        if (plugin == null) {
            // quite！！！
            return;
        }

        // 本地插件, 不做处理
        if (plugin.isNativePlugin()) {
        } else {
            lock.writeLock().lock();
            try {
                if (plugin.isSeparateClassLoaderNeeded()) {
                    // 关闭类加载器
                    String cacheKey = ((Plugin)plugin).getPluginUniqueCode();
                    List<URLClassLoader> classLoaders = separateClassLoaderHandler.remove(cacheKey);
                    classLoaders.forEach(clm::closeClassLoader);
                } else {
                    // 关闭缓存中的类加载器
                    removeCachedClassLoader(plugin);
                }
            } finally {
                lock.writeLock().unlock();
            }
        }
    }

    /**
     * 创建一个新的类加载器
     *
     * @param plugin
     * @return
     */
    private URLClassLoader createClassLoader(PluginInterface plugin) {
        String pluginUniqueCode = ((Plugin)plugin).getPluginUniqueCode();

        List<File> mainJars = FileUtil.loopFiles(new File(plugin.getPluginDirectoryString()), 1,
            f -> f.getName().toLowerCase().endsWith(".jar"));

        List<File> libJars =
            plugin.getLibraries().stream().map(File::new).filter(File::exists).collect(Collectors.toList());

        URL[] urls = Stream.concat(mainJars.stream(), libJars.stream()).map(URLUtil::getURL).toArray(URL[]::new);

        String[] patterns = parentClassLoaderPatternMap.get(plugin);
        // 该插件中有需要从父类加载器加载的类
        if (patterns != null) {
            return new KettleSelectiveParentFirstClassLoader(urls, getClass().getClassLoader(), pluginUniqueCode,
                patterns);
        } else {
            return new KettleURLClassLoader(urls, getClass().getClassLoader(), pluginUniqueCode);
        }
    }

    /**
     * 关闭类加载器
     *
     * @param classLoader
     */
    private void closeClassLoader(URLClassLoader classLoader) {
        try {
            if (classLoader != null) {
                classLoader.close();
                if (classLoader instanceof KettleURLClassLoader) {
                    ((KettleURLClassLoader)classLoader).closeClassLoader();
                }
            }
        } catch (IOException e) {
        }
    }

    /**
     * 获取类加载器缓存
     *
     * @return
     */
    private Map<String, URLClassLoader> getClassLoaderCache(CacheType cacheType) {
        return classLoaderMap.get(cacheType.name());
    }

    /**
     * 根据插件获取类加载器，优先从缓存中获取，如缓存中不存在则创建新的类加载器.
     *
     * @param
     * @param
     * @return
     */
    private URLClassLoader getCachedClassLoader(PluginInterface plugin) {
        String cacheKey = null;

        // userDefined
        cacheKey = plugin.getClassLoaderCfg().getUserDefinedGroup();
        if (StrUtil.isNotBlank(cacheKey)) {
            URLClassLoader classLoader = getClassLoaderCache(CacheType.USER_DEFINED).get(cacheKey);
            if (classLoader == null) {
                classLoader = createClassLoader(plugin);
                getClassLoaderCache(CacheType.USER_DEFINED).put(cacheKey, classLoader);
            }
            return classLoader;
        }

        // folderBased
        Boolean folderBased = plugin.getClassLoaderCfg().getFolderBasedGroup();
        if (folderBased) {
            cacheKey = plugin.getPluginDirectory().toString();
            URLClassLoader classLoader = getClassLoaderCache(CacheType.FOLDER_BASED).get(cacheKey);
            if (classLoader == null) {
                classLoader = createClassLoader(plugin);
                getClassLoaderCache(CacheType.FOLDER_BASED).put(cacheKey, classLoader);
            }
            return classLoader;
        }

        // other
        cacheKey = ((Plugin)plugin).getPluginUniqueCode();
        URLClassLoader classLoader = getClassLoaderCache(CacheType.UNIQUE_CODE).get(cacheKey);
        if (classLoader == null) {
            classLoader = createClassLoader(plugin);
            getClassLoaderCache(CacheType.UNIQUE_CODE).put(cacheKey, classLoader);
        }
        return classLoader;
    }

    private void removeCachedClassLoader(PluginInterface plugin) {
        String cacheKey = null;
        URLClassLoader classLoader = null;

        // userDefined
        cacheKey = plugin.getClassLoaderCfg().getUserDefinedGroup();
        if (StrUtil.isNotBlank(cacheKey)) {
            classLoader = getClassLoaderCache(CacheType.USER_DEFINED).remove(cacheKey);
            closeClassLoader(classLoader);
        }

        // folderBased
        Boolean folderBased = plugin.getClassLoaderCfg().getFolderBasedGroup();
        if (folderBased) {
            cacheKey = plugin.getPluginDirectory().toString();
            classLoader = getClassLoaderCache(CacheType.FOLDER_BASED).remove(cacheKey);
            closeClassLoader(classLoader);
        }

        // other
        cacheKey = ((Plugin)plugin).getPluginUniqueCode();
        classLoader = getClassLoaderCache(CacheType.UNIQUE_CODE).remove(cacheKey);
        closeClassLoader(classLoader);
    }

}
